from langchain import FAISS
from langchain.document_loaders import TextLoader, UnstructuredFileLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
import sentence_transformers

from 压铸机5.MyFAISS import MyFAISS

# 导入文本
loader = UnstructuredFileLoader("压铸机2.txt", mode="elements")  # 使用 UnstructuredFileLoader 加载器以元素模式加载文件
text_splitor = CharacterTextSplitter()  # 使用 CharacterTextSplitter 来分割文件中的文本
docs = loader.load_and_split(text_splitor)  # 加载文件并进行文本分割

embedding_model_dict = {
    "text2vec": "C:\\Work\\llm\\text2vec-large-chinese",
    "m3e-base": "C:\\Work\\llm\\m3e-base"
}

EMBEDDING_MODEL = "m3e-base"
# 初始化 hugginFace 的 embeddings 对象
embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[EMBEDDING_MODEL])
embeddings.client = sentence_transformers.SentenceTransformer(embeddings.model_name, device='cuda')

# 初始化加载器
db = MyFAISS.from_documents(docs, embeddings)
# 匹配后单段上下文长度
db.chunk_size = 250
db.chunk_conent = True
# 知识检索内容相关度 Score, 数值范围约为0-1100，如果为0，则不生效，经测试设置为小于500时，匹配结果更精准
db.score_threshold = 500

# # 持久化
db.save_local("./faiss/casting2")
